/*
 * Copyright 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

@file:JvmName("DataStoreDelegateKt") // Workaround for b/313964643

package androidx.datastore

import android.annotation.SuppressLint
import android.content.Context
import android.os.Build
import androidx.annotation.GuardedBy
import androidx.annotation.RequiresApi
import androidx.datastore.core.DataMigration
import androidx.datastore.core.DataStore
import androidx.datastore.core.DataStoreFactory
import androidx.datastore.core.Serializer
import androidx.datastore.core.deviceProtectedDataStoreFile
import androidx.datastore.core.handlers.ReplaceFileCorruptionHandler
import androidx.datastore.core.okio.OkioSerializer
import androidx.datastore.core.okio.OkioStorage
import kotlin.properties.ReadOnlyProperty
import kotlin.reflect.KProperty
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import okio.BufferedSink
import okio.BufferedSource
import okio.FileSystem
import okio.Path.Companion.toPath

/**
 * Creates a property delegate for a single process DataStore. This should only be called once in a
 * file (at the top level), and all usages of the DataStore should use a reference the same
 * Instance. The receiver type for the property delegate must be an instance of [Context].
 *
 * This should only be used from a single application in a single classloader in a single process.
 *
 * Example usage:
 * ```
 * val Context.myDataStore by dataStore("filename", serializer)
 *
 * class SomeClass(val context: Context) {
 *    suspend fun update() = context.myDataStore.updateData {...}
 * }
 * ```
 *
 * @param fileName the filename relative to Context.applicationContext.filesDir that DataStore acts
 *   on. The File is obtained from [dataStoreFile]. It is created in the "/datastore" subdirectory,
 *   in the user encrypted storage.
 * @param serializer The serializer for `T`.
 * @param corruptionHandler The corruptionHandler is invoked if DataStore encounters a
 *   [androidx.datastore.core.CorruptionException] when attempting to read data.
 *   CorruptionExceptions are thrown by serializers when data can not be de-serialized.
 * @param produceMigrations produce the migrations. The ApplicationContext is passed in to these
 *   callbacks as a parameter. DataMigrations are run before any access to data can occur. Each
 *   producer and migration may be run more than once whether or not it already succeeded
 *   (potentially because another migration failed or a write to disk failed.)
 * @param scope The scope in which IO operations and transform functions will execute.
 * @return a property delegate that manages a datastore as a singleton.
 */
@Suppress("MissingJvmstatic")
public fun <T> dataStore(
    fileName: String,
    serializer: Serializer<T>,
    corruptionHandler: ReplaceFileCorruptionHandler<T>? = null,
    produceMigrations: (Context) -> List<DataMigration<T>> = { listOf() },
    scope: CoroutineScope = CoroutineScope(Dispatchers.IO + SupervisorJob()),
): ReadOnlyProperty<Context, DataStore<T>> {
    return DataStoreSingletonDelegate(
        fileName,
        OkioSerializerWrapper(serializer),
        corruptionHandler,
        produceMigrations,
        scope,
        createInDeviceProtectedStorage = false,
    )
}

/**
 * Creates a property delegate for a single process DataStore to be used in direct boot.
 *
 * @param fileName the filename relative to Context.applicationContext.filesDir that DataStore acts
 *   on. The File is obtained from [dataStoreFile]. It is created in the "/datastore" subdirectory,
 *   in the user encrypted storage.
 * @param serializer The serializer for `T`.
 * @param corruptionHandler The corruptionHandler is invoked if DataStore encounters a
 *   [androidx.datastore.core.CorruptionException] when attempting to read data.
 *   CorruptionExceptions are thrown by serializers when data can not be de-serialized.
 * @param produceMigrations produce the migrations. The ApplicationContext is passed in to these
 *   callbacks as a parameter. DataMigrations are run before any access to data can occur. Each
 *   producer and migration may be run more than once whether or not it already succeeded
 *   (potentially because another migration failed or a write to disk failed.)
 * @param scope The scope in which IO operations and transform functions will execute.
 * @return a property delegate that manages a datastore as a singleton.
 */
@Suppress("MissingJvmstatic")
@RequiresApi(Build.VERSION_CODES.N)
public fun <T> deviceProtectedDataStore(
    fileName: String,
    serializer: Serializer<T>,
    corruptionHandler: ReplaceFileCorruptionHandler<T>? = null,
    produceMigrations: (Context) -> List<DataMigration<T>> = { listOf() },
    scope: CoroutineScope = CoroutineScope(Dispatchers.IO + SupervisorJob()),
): ReadOnlyProperty<Context, DataStore<T>> {
    return DataStoreSingletonDelegate(
        fileName,
        OkioSerializerWrapper(serializer),
        corruptionHandler,
        produceMigrations,
        scope,
        createInDeviceProtectedStorage = true,
    )
}

/** Delegate class to manage DataStore as a singleton. */
internal class DataStoreSingletonDelegate<T>
internal constructor(
    private val fileName: String,
    private val serializer: OkioSerializer<T>,
    private val corruptionHandler: ReplaceFileCorruptionHandler<T>?,
    private val produceMigrations: (Context) -> List<DataMigration<T>>,
    private val scope: CoroutineScope,
    private val createInDeviceProtectedStorage: Boolean,
) : ReadOnlyProperty<Context, DataStore<T>> {

    private val lock = Any()

    @GuardedBy("lock") @Volatile private var INSTANCE: DataStore<T>? = null

    /**
     * Gets the instance of the DataStore.
     *
     * This function is not guarded by a [RequiresApi] annotation as only [deviceProtectedDataStore]
     * will set [createInDeviceProtectedStorage] to true, resulting in a call to
     * [deviceProtectedDataStoreFile]. Since [deviceProtectedDataStore] is guarded by @RequiresApi,
     * we do not set an API requirement on [getValue] so [deviceProtectedDataStore] and [dataStore]
     * property delegates can share the function.
     *
     * @param thisRef must be an instance of [Context]
     * @param property not used
     */
    @SuppressLint("NewApi")
    override fun getValue(thisRef: Context, property: KProperty<*>): DataStore<T> {
        return INSTANCE
            ?: synchronized(lock) {
                if (INSTANCE == null) {
                    val applicationContext = thisRef.applicationContext
                    INSTANCE =
                        DataStoreFactory.create(
                            storage =
                                OkioStorage(FileSystem.SYSTEM, serializer) {
                                    if (createInDeviceProtectedStorage) {
                                            // Should not be able to reach here with lower SDK level
                                            check(Build.VERSION.SDK_INT >= Build.VERSION_CODES.N)
                                            applicationContext.deviceProtectedDataStoreFile(
                                                fileName
                                            )
                                        } else {
                                            // Regardless of whether thisRef is a
                                            // DeviceProtectedStorageContext, the applicationContext
                                            // will always provide a User-Encrypted Storage Context.
                                            applicationContext.dataStoreFile(fileName)
                                        }
                                        .absolutePath
                                        .toPath()
                                },
                            corruptionHandler = corruptionHandler,
                            migrations = produceMigrations(applicationContext),
                            scope = scope,
                        )
                }
                INSTANCE!!
            }
    }
}

internal class OkioSerializerWrapper<T>(private val delegate: Serializer<T>) : OkioSerializer<T> {
    override val defaultValue: T
        get() = delegate.defaultValue

    override suspend fun readFrom(source: BufferedSource): T {
        return delegate.readFrom(source.inputStream())
    }

    override suspend fun writeTo(t: T, sink: BufferedSink) {
        delegate.writeTo(t, sink.outputStream())
    }
}
